﻿// ESP.cpp : 定义静态库的函数。
//

#include "pch.h"
#include "framework.h"

#include <vector>
#include <Windows.h>

#include "ESP.h"
#include "GetESP.h"
#include "GetVolume.h"
#include "MyAlloc.h"

std::vector<char> AvaibleLetterList;
std::vector<_PartitionInfoW*> _PIWL;

//由于C++回收机制，直接向MountedEspLetter/UnmountedEspLetter复制的结果是MountedEspLetter/UnmountedEspLetter被回收，所以我们需要一个全局变量作中转站“EspLetter”（全局变量不会被回收）。
// 赋值时，先往EspLetter赋值，再通过赋值到MountedEspLetter/UnmountedEspLetter，方式曲折但有效。
std::vector<char> EspLetter;

std::vector<char> GetAvaibleLetter();
bool Init();

BOOL MountDriveW(char Letter, std::wstring PartionDrivePath);
BOOL UnmountDriveW(char Letter, std::wstring PartionDrivePath);
BOOL UnmountDriveW(char Letter);


/// <summary>
/// 获取电脑中所有的ESP分区
/// </summary>
/// <returns>成功返回std::vector，失败返回NULL，可用GetLastError查看详细信息</returns>
std::vector<_PartitionInfoW*> Esp_GetAllEspPartitionW() {
	return GetEspPartitionInfo();
}


/// <summary>
/// 挂载所有EFI分区，已挂载的分区不会重复挂载
/// </summary>
/// <param name="MountedEspLetter">指向用于接收挂载的EFI分区盘符的指针</param>
/// <returns>成功返回1，失败返回0，可使用GetLastError获取错误信息</returns>
BOOL Esp_MountEfiPartition(_Out_ std::vector<char>* MountedEspLetter) {
	//std::vector<_PartitionInfoW*> _PIWL;
	//std::vector<char> MountedEspLetterA;

	//开始之前，先清除EspLetter，
	EspLetter.clear();

	//初始化
	bool Status = Init();

	//初始化失败，保留错误信息，返回
	if (Status == false) {
		return 0;
	}

	//检查盘符是否足够挂载ESP分区，若不能，报错ERROR_MOUNTESP_NO_MORE_AVAIBLE_LETTER，返回
	if (_PIWL.size() > AvaibleLetterList.size()) {
		MyFreeByLessIndex(4);
		SetLastError(ERROR_MOUNTESP_NO_MORE_AVAIBLE_LETTER);
		return 0;
	}

	//挂载ESP分区：
	for (int i = 0; i < _PIWL.size(); i++) {
		
		//若ESP分区已挂载，则跳过
		if (_PIWL[i]->HasPartitionLetter) {
			continue;
		}

		//挂载分区
		if (MountDriveW(AvaibleLetterList[i], _PIWL[i]->PartionDrivePathW) == 0) {
			MyFreeByLessIndex(4);
			return 0;
		}

		//将已挂载的ESP分区加入到EspLetter
		EspLetter.push_back(AvaibleLetterList[i]);
	}

	//正式赋值
	*MountedEspLetter = EspLetter;

	MyFreeByLessIndex(4);
	return 1;
}

/// <summary>
/// 挂载指定硬盘编号的EFI分区，已挂载的不会重新挂载
/// </summary>
/// <param name="DiskNumber">硬盘号，从0开始</param>
/// <param name="MountedEspLetter">指向用于接收挂载的EFI分区盘符的指针</param>
/// <returns>成功返回1，失败返回0，可使用GetLastError获取错误信息</returns>
BOOL Esp_MountEfiPartition(_In_ int DiskNumber, _Out_ std::vector<char>* MountedEspLetter) {
	std::vector<_PartitionInfoW*> _Disk_ESP;
	//std::vector<char> MountedEspLetter;
	
	//开始之前，先清除EspLetter，
	EspLetter.clear();

	//初始化
	bool Status = Init();

	//初始化失败，保留错误信息，返回
	if (Status == false) {
		return 0;
	}

	//通过遍历的方法获取指定磁盘
	for (int i = 0; i < _PIWL.size(); i++) {
		//若ESP分区已挂载，或者不是指定硬盘，则跳过
		if (_PIWL[i]->HasPartitionLetter || _PIWL[i]->DiskNumber != DiskNumber) {
			continue;
		}

		//挂载分区
		if (MountDriveW(AvaibleLetterList[i], _PIWL[i]->PartionDrivePathW) == 0) {
			MyFreeByLessIndex(4);
			return 0;
		}

		//将已挂载的ESP分区加入到MountedEspLetter
		EspLetter.push_back(AvaibleLetterList[i]);
	}

	*MountedEspLetter = EspLetter;

	MyFreeByLessIndex(4);
	return 1;
}

/// <summary>
/// 卸载所有EFI分区，已卸载的不会重复卸载
/// </summary>
/// <param name="UnmountedEspLetter">指向用于接收卸载的EFI分区盘符的指针</param>
/// <returns>成功返回1,失败返回0，可使用GetLastError获取错误信息</returns>
BOOL Esp_UnmountEfiPartition(_Out_ std::vector<char>* UnmountedEspLetter) {
	//开始之前，先清除EspLetter，
	EspLetter.clear();

	//初始化
	bool Status = Init();

	//初始化失败，保留错误信息，返回
	if (Status == false) {
		return 0;
	}

	//通过遍历方法卸载ESP分区
	for (int i = 0; i < _PIWL.size(); i++) {
		//若分区未挂载，则跳过
		if (!_PIWL[i]->HasPartitionLetter) {
			continue;
		}

		//卸载分区
		BOOL Status = UnmountDriveW(_PIWL[i]->PartitionLetter, _PIWL[i]->PartionDrivePathW);

		//挂载失败，报错退出
		if (Status == 0) {
			MyFreeByLessIndex(4);
			return 0;
		}
		
		//将卸载后的盘符推到UnmountedEspLetter
		EspLetter.push_back(_PIWL[i]->PartitionLetter);
	}

	*UnmountedEspLetter = EspLetter;

	MyFreeByLessIndex(4);
	return 1;
}

/// <summary>
/// 卸载指定EFI分区，已卸载的不会重复卸载
/// </summary>
/// <param name="Letter">盘符，例如"C","D"，字母必须大写</param>
/// <returns>成功返回1，失败返回0，可使用GetLastError获取错误信息</returns>
BOOL Esp_UnmountEfiPartition(_In_ char Letter) {
	std::vector<_PartitionInfoW*> _PIWL_A;
	std::vector<char> MountedEspLetter;

	_PIWL_A = Esp_GetAllEspPartitionW(); //获取所有ESP分区

	//获取失败，报错
	if (_PIWL_A == (std::vector<_PartitionInfoW*>)NULL) {
		return 0;
	}

	for (int i = 0; i < _PIWL_A.size(); i++) {
		//找到则卸载该磁盘
		if (_PIWL_A[i]->PartitionLetter == Letter) {
			UnmountDriveW(Letter);
			MountedEspLetter.push_back(Letter);
		}
	}

	return 1;
	
}

/// <summary>
/// 卸载指定硬盘编号的EFI分区，已卸载的不会重新卸载
/// </summary>
/// <param name="DiskNumber">硬盘号，从0开始</param>
/// <param name="UnmountedEspLetter">指向用于接收卸载的EFI分区盘符的指针</param>
/// <returns>成功返回1,失败返回0，可使用GetLastError获取错误信息/returns>
BOOL Esp_UnmountEfiPartition(_In_ int DiskNumber,_Out_ std::vector<char>* UnmountedEspLetter) {
	//开始之前，先清除EspLetter，
	EspLetter.clear();

	std::vector<_PartitionInfoW*> _Disk_ESP;
	//std::vector<char> UnmountedEspLetter;

	//初始化
	bool Status = Init();

	//初始化失败，保留错误信息，返回
	if (Status == false) {
		return 0;
	}

	for (int i = 0; i < _PIWL.size(); i++) {
		if (_PIWL[i]->DiskNumber == DiskNumber) {
			_Disk_ESP.push_back(_PIWL[i]);
		}
	}

	//卸载分区
	for (int i = 0; i < _Disk_ESP.size(); i++) {
		//若ESP分区未挂载，则跳过
		if (!_Disk_ESP[i]->HasPartitionLetter) {
			continue;
		}

		//卸载分区
		if (UnmountDriveW(_Disk_ESP[i]->PartitionLetter) == 0 ){
			MyFreeByLessIndex(4);
			return 0;
		}

		//将已挂载的ESP分区加入到MountedEspLetter
		EspLetter.push_back(_Disk_ESP[i]->PartitionLetter);
	}

	*UnmountedEspLetter = EspLetter;

	MyFreeByLessIndex(4);
	return 1;
}

/// <summary>
/// 获取可用盘符
/// </summary>
/// <returns></returns>
std::vector<char> GetAvaibleLetter() {
	std::vector<char> AvaibleLetterList;
	
	//创建完全盘符列表
	std::vector<char> AllLetter;
	for (char i = 'A'; i <= 'Z'; i++) {
		AllLetter.push_back(i);
	}

	//获取计算机中已使用的盘符
	std::vector<char> MountedLetterList = GetVolumeLetter();

	for (int i = 0; i < AllLetter.size(); i++) {
		//当查找不到时，说明该盘符未使用
		if(std::find(MountedLetterList.begin(), MountedLetterList.end(), AllLetter[i]) == MountedLetterList.end())
		{
			AvaibleLetterList.push_back(AllLetter[i]);
		}
	}

	//如果没有多于盘符，报错ERROR_GETESPINFO_MO_MORE_AVAIBLE_LETTER，返回
	if (AvaibleLetterList.size() == 0) {
		SetLastError(ERROR_GETESPINFO_NO_MORE_AVAIBLE_LETTER);
		return (std::vector<char>)-1;
	}

	return AvaibleLetterList;
}

bool Init() {
	AvaibleLetterList.clear();
	_PIWL.clear();
	try {
		//获取ESP分区
		std::vector<_PartitionInfoW*> _PIWLW = GetEspPartitionInfo();

		//获取失败，保留错误信息，返回
		if (_PIWLW == (std::vector<_PartitionInfoW*>)NULL) {
			return false;
		}

		//获取可用盘符
		std::vector<char> AvaibleLetterListW = GetAvaibleLetter();
		//获取失败，保留错误信息，返回
		if (AvaibleLetterListW == (std::vector<char>)NULL) {
			return false;
		}

		//_PIWL = _PIWLW;
		//AvaibleLetterList = AvaibleLetterListW;

		for (int i = 0; i < _PIWLW.size(); i++) {
			_PIWL.push_back(_PIWLW[i]);
		}

		for (int i = 0; i < AvaibleLetterListW.size(); i++) {
			AvaibleLetterList.push_back(AvaibleLetterListW[i]);
		}

		return true;
	}
	catch(const char* msg){
		return false;
	}
}

/// <summary>
/// 挂载分区
/// </summary>
/// <param name="Letter">盘符，例如"C","D"</param>
/// <param name="PartionDrivePath">分区路径，如“\\Device\\Harddisk0\\Partition1”</param>
/// <returns>成功返回非0值，失败返回0，可用GetLastError查询错误信息</returns>
BOOL MountDriveW(char Letter,std::wstring PartionDrivePath) {
	//拼接字符串，拼接后字符： X: （X为盘符）
	std::wstring ws;
	ws += Letter;
	ws += L":";

	//挂载分区
	return DefineDosDevice(DDD_RAW_TARGET_PATH, ws.c_str(), PartionDrivePath.c_str());
	
}

/// <summary>
/// 卸载分区
/// </summary>
/// <param name="Letter">盘符，例如"C","D"</param>
/// <param name="PartionDrivePath">分区路径，如“\\Device\\Harddisk0\\Partition1”</param>
/// <returns>成功返回非0值，失败返回0，可用GetLastError查询错误信息</returns>
BOOL UnmountDriveW(char Letter, std::wstring PartionDrivePath) {
	//拼接字符串，拼接后字符： X: （X为盘符）
	std::wstring ws;
	ws += Letter;
	ws += L":";

	//卸载分区
	return DefineDosDevice(DDD_REMOVE_DEFINITION | DDD_EXACT_MATCH_ON_REMOVE | DDD_RAW_TARGET_PATH, ws.c_str(), PartionDrivePath.c_str());
}

/// <summary>
/// 卸载分区
/// </summary>
/// <param name="Letter">盘符，例如"C","D"，字母必须大写</param>
/// <returns>成功返回非0值，失败返回0，可用GetLastError查询错误信息</returns>
BOOL UnmountDriveW(char Letter) {
	//拼接字符串，拼接后字符： X: （X为盘符）
	std::wstring ws;
	ws += Letter;
	ws += L":";

	//卸载分区
	return DefineDosDevice(DDD_REMOVE_DEFINITION, ws.c_str(), nullptr);

}

